import math
import random
import pandas as pd

def predict_hsa(egfr, age, bmi, is_female=False, is_black=False,
                add_noise=False, sd=None, return_unit='g/dL'):
    """
    使用多变量回归模型预测人血清白蛋白（HSA）浓度，并可进行单位转换，
    同时可根据模型残差标准差添加噪声（epsilon）。

    默认残差 SD：
      • eGFR < 120 按 CKD 分期：
        - CKD 5 (egfr<15): SD = 0.42
        - CKD 4 (15≤egfr<30): SD = 0.38
        - CKD 3 (30≤egfr<60): SD = 0.35
        - CKD 2 (60≤egfr<90): SD = 0.32
        - CKD 1 (90≤egfr<120): SD = 0.29
      • eGFR ≥ 120 按年龄分组：
        - age < 35: SD = 0.30
        - 35 ≤ age < 45: SD = 0.32
        - 45 ≤ age < 55: SD = 0.35
        - age ≥ 55: SD = 0.50

    参数:
        egfr (float): 估计肾小球滤过率 (mL/min/1.73m²)
        age (float): 年龄 (岁)
        bmi (float): 体重指数 (kg/m²)
        is_female (bool): 是否为女性
        is_black (bool): 是否为黑人
        add_noise (bool): 是否添加随机噪声
        sd (float or dict): 覆盖默认 SD；若为 dict，可包含 'lt120'/'gte120' 键
        return_unit (str): 'g/dL' 或 'g/L'
    返回:
        float: 预测的 HSA 浓度
    """
    # 归一化
    x1, x2, x3 = egfr/90.0, age/50.0, bmi/30.0

    # 基础预测
    if egfr < 120:
        f_fac = 0.967 if is_female else 1.0
        b_fac = 0.973 if is_black else 1.0
        hsa_base = (
            0.824*x1 - 0.378*x1**2
            - 1.046*x2 + 0.797*x2**2 - 0.223*x2**3
            + 0.896*x3 - 1.008*x3**2 + 0.247*x3**3
            + 4.227
        )
        hsa = hsa_base * f_fac * b_fac
    else:
        f_fac = 0.941 if is_female else 1.0
        b_fac = 0.964 if is_black else 1.0
        hsa_base = (
            -0.213*x2 - 0.295*x3 - 0.215*x3**2
            + 0.067*x3**3 + 5.063
        )
        hsa = hsa_base * f_fac * b_fac

    # 确定 SD
    if sd is None:
        if egfr < 120:
            if egfr < 15:
                noise_sd = 0.42
            elif egfr < 30:
                noise_sd = 0.38
            elif egfr < 60:
                noise_sd = 0.35
            elif egfr < 90:
                noise_sd = 0.30
                
            else:
                noise_sd = 0.29
        else:
            if age < 35:
                noise_sd = 0.30
            elif age < 45:
                noise_sd = 0.32
            elif age < 55:
                noise_sd = 0.35
            else:
                noise_sd = 0.50
    else:
        if isinstance(sd, dict):
            if egfr < 120 and 'lt120' in sd:
                noise_sd = sd['lt120']
            elif egfr >= 120 and 'gte120' in sd:
                noise_sd = sd['gte120']
            else:
                noise_sd = list(sd.values())[0]
        else:
            noise_sd = sd

    # 添加噪声
    if add_noise:
        hsa += random.gauss(0, noise_sd)

    # 单位转换
    return hsa*10.0 if return_unit=='g/L' else hsa

def generate_blood_parameters_from_df(df):
    """
    批量生成血液参数，输入人口学DataFrame，输出带id的血液参数DataFrame。
    """
    results = []
    for idx, row in df.iterrows():
        person_id = row.get('id', idx)
        egfr = row.get('egfr', 100)  # 默认值100，可根据实际数据调整
        age = row.get('age', 40)
        bmi = row.get('bmi', 22)
        hsa = predict_hsa(egfr, age, bmi)
        results.append({
            'id': person_id,
            'Blood_HSA': hsa
        })
    return pd.DataFrame(results)

# 示例
if __name__ == "__main__":
    import pandas as pd
    # 构造测试数据
    test_df = pd.DataFrame({
        'id': [1, 2],
        'age': [30, 40],
        'bmi': [22, 24],
        'egfr': [90, 110]
    })
    result = generate_blood_parameters_from_df(test_df)
    print(result)
    print(predict_hsa(12, 60, 25, is_female=True, add_noise=True))      # CKD5 期
    print(predict_hsa(130, 60, 22, add_noise=True, return_unit='g/L'))  # 高 eGFR 组
